import yolov3_model
from utils import *
import matplotlib.pyplot as plt
path = ''
model = yolov3_model.model
model.eval()
def predict(path):
    img0 = cv2.imread(path, -1)
    imgshape = img0.shape[0:2]
    anchor = []
    img = transImg(img0)
    img = paddle.unsqueeze(img, 0)
    img = img.astype(paddle.float32)

    out = model(img)
    box, score, classes = eval(out, anchor, imgshape, 4)
    for index, c in enumerate(classes,0):
        _box = box[index]
        _score = score[index]
        top, left, bottom, right = box
        top = np.max(0,np.floor(top+0.5)).astype('int32')
        left = np.max(0, np.floor(left + 0.5)).astype('int32')
        bottom = np.min(imgshape[1], np.floor(bottom + 0.5)).astype('int32')
        right = np.min(imgshape[0], np.floor(right + 0.5)).astype('int32')
        cv2.rectangle(img0, (top, left), (bottom, right), (200, 0, 0), 1)
        cv2.putText(img0, 'class: '+str(c), (top+3, left+3), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 255), 2)
    plt.imshow(img0)
    plt.show()


if __name__ == '__main__':
    predict(path)